package com.hzsparrow.ai.service;

import java.util.List;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.template.st.StTemplateRenderer;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.hzsparrow.ai.config.PromptConfiguration;
import com.hzsparrow.ai.config.AIModelInitService;

import reactor.core.publisher.Flux;

/**
 * 聊天服务类
 */
@Service
public class ChatService {

    private final PromptConfiguration promptConfig;
    private final AIModelInitService aiModelInitService;

    @Autowired
    public ChatService(PromptConfiguration promptConfig, AIModelInitService aiModelInitService) {
        this.promptConfig = promptConfig;
        this.aiModelInitService = aiModelInitService;
    }

    /**
     * 处理聊天消息
     * 
     * @param message 用户消息
     * @return 流式AI响应
     */
    public Flux<String> processChat(String message) {
        return processChat(message, null);
    }
    
    /**
     * 处理聊天消息，支持指定知识库
     * 
     * @param message 用户消息
     * @param vectorIndexName 向量索引名称，为null时使用默认索引
     * @return 流式AI响应
     */
    public Flux<String> processChat(String message, String vectorIndexName) {
        
        // 获取当前默认的聊天模型
        ChatModel chatModel = aiModelInitService.getDefaultChatModel();
        if (chatModel == null) {
            return Flux.error(new IllegalStateException("默认聊天模型未配置"));
        }
        
        // 获取向量存储
        VectorStore vectorStore = aiModelInitService.getDefaultVectorStore();
        
        if (vectorStore == null) {
            return Flux.error(new IllegalStateException("向量存储未配置"));
        }
        
        // 创建ChatClient构建器
        ChatClient.Builder chatClientBuilder = ChatClient.builder(chatModel);

        // 如果指定了知识库，则进入知识库回答判断流程
        if (vectorIndexName != null && !vectorIndexName.isEmpty()) {
            // 创建SearchRequest构建器
            FilterExpressionBuilder b = new FilterExpressionBuilder();
            SearchRequest.Builder searchRequestBuilder = SearchRequest.builder()
            .query(message)
            .similarityThreshold(0.6d)
            .topK(6)
            .filterExpression(b.eq("index_name", vectorIndexName).build());

            SearchRequest searchRequest = searchRequestBuilder.build();

            // 首先尝试从向量库检索相关文档
            List<Document> relevantDocs = vectorStore.similaritySearch(searchRequest);
            
            // 判断是否找到相关文档
            if (!relevantDocs.isEmpty()) {
                // 找到相关文档，使用RAG方式回答
                System.out.println("找到相关文档，使用RAG方式回答，文档数量: " + relevantDocs.size());
                
                // 创建QuestionAnswerAdvisor用于RAG
                PromptTemplate customPromptTemplate = PromptTemplate.builder()
                        .renderer(StTemplateRenderer.builder().startDelimiterToken('{').endDelimiterToken('}').build())
                        .template(promptConfig.getVectorChatPrompt())
                        .build();
                QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                        .searchRequest(searchRequest)
                        .promptTemplate(customPromptTemplate)
                        .build();
                
                // 返回RAG方式的回答
                return chatClientBuilder.build().prompt(message)
                        .advisors(questionAnswerAdvisor)
                        .stream()
                        .content();
            }
        }

        // 未指定知识库或未找到相关文档，使用普通方式回答
        System.out.println("未指定知识库或未找到相关文档，使用普通方式回答");
        
        // 创建包含系统提示和用户提示的Prompt
        Prompt prompt = new Prompt(
                new SystemMessage(promptConfig.getGeneralChatPrompt()),
                new UserMessage(message)
        );
        
        // 返回普通方式的回答
        return chatClientBuilder.build().prompt(prompt)
                .stream()
                .content();
    }
    
} 